import os
import re
import pandas as pd
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment, Border, Side, PatternFill

# 尝试导入fitz，处理可能的导入错误
try:
    import fitz  # PyMuPDF
except ImportError as e:
    print(f"导入PyMuPDF时出错: {e}")
    print("请尝试重新安装PyMuPDF: pip uninstall pymupdf fitz && pip install pymupdf")
    print("或者在项目根目录创建static文件夹: mkdir static")
    exit(1)
except RuntimeError as e:
    print(f"PyMuPDF运行时错误: {e}")
    print("请在项目根目录创建static文件夹: mkdir static")
    exit(1)


def extract_financial_tables(pdf_path):
    """从PDF文件中提取财务报表数据，使用PyMuPDF（更简单可靠）"""
    print(f"正在处理PDF文件: {pdf_path}")

    try:
        # 打开PDF文件
        doc = fitz.open(pdf_path)
    except Exception as e:
        print(f"打开PDF文件时出错: {e}")
        return {
            "balance_sheet": [],
            "income_statement": [],
            "cash_flow": []
        }

    # 用于存储提取的表格数据
    balance_sheet_data = []
    income_statement_data = []
    cash_flow_data = []

    # 定义关键词，用于识别各类报表
    balance_sheet_keywords = ["合并资产负债表", "资产负债表"]
    income_statement_keywords = ["合并利润表", "利润表"]
    cash_flow_keywords = ["合并现金流量表", "现金流量表"]

    # 查找报表所在页面
    balance_sheet_pages = []
    income_statement_pages = []
    cash_flow_pages = []

    # 查找报表页码
    for page_num in range(len(doc)):
        page = doc[page_num]
        text = page.get_text()

        # 检查是否包含各报表关键词
        if any(keyword in text for keyword in balance_sheet_keywords):
            print(f"在第 {page_num + 1} 页找到资产负债表")
            balance_sheet_pages.append(page_num)

        if any(keyword in text for keyword in income_statement_keywords):
            print(f"在第 {page_num + 1} 页找到利润表")
            income_statement_pages.append(page_num)

        if any(keyword in text for keyword in cash_flow_keywords):
            print(f"在第 {page_num + 1} 页找到现金流量表")
            cash_flow_pages.append(page_num)

    # 提取资产负债表数据
    if balance_sheet_pages:
        balance_sheet_data = extract_table_from_pages(doc, balance_sheet_pages, "balance_sheet")

    # 提取利润表数据
    if income_statement_pages:
        income_statement_data = extract_table_from_pages(doc, income_statement_pages, "income_statement")

    # 提取现金流量表数据
    if cash_flow_pages:
        cash_flow_data = extract_table_from_pages(doc, cash_flow_pages, "cash_flow")

    # 关闭PDF文件
    doc.close()

    return {
        "balance_sheet": balance_sheet_data,
        "income_statement": income_statement_data,
        "cash_flow": cash_flow_data
    }


def extract_table_from_pages(doc, page_numbers, table_type):
    """从指定页面提取财务表格数据"""
    table_data = []

    # 计算结束页面 - 每种表通常会占据多页
    start_page = page_numbers[0]

    # 根据表格类型设置提取参数
    if table_type == "balance_sheet":
        items_column_patterns = [r"资产", r"负债", r"所有者权益", r"货币资金", r"应收账款", r"存货",
                                 r"流动资产", r"非流动资产", r"资产总计", r"应付账款", r"流动负债",
                                 r"非流动负债", r"负债总计", r"实收资本", r"股本", r"资本公积",
                                 r"盈余公积", r"未分配利润"]
        value_patterns = [r"\d{1,3}(,\d{3})*(\.\d+)?", r"-\d{1,3}(,\d{3})*(\.\d+)?"]
        max_pages = 4  # 资产负债表通常不超过4页
    elif table_type == "income_statement":
        items_column_patterns = [r"营业收入", r"营业成本", r"营业利润", r"利润总额", r"净利润",
                                 r"销售费用", r"管理费用", r"财务费用", r"研发费用", r"所得税费用"]
        value_patterns = [r"\d{1,3}(,\d{3})*(\.\d+)?", r"-\d{1,3}(,\d{3})*(\.\d+)?"]
        max_pages = 3  # 利润表通常不超过3页
    elif table_type == "cash_flow":
        items_column_patterns = [r"经营活动", r"投资活动", r"筹资活动", r"现金流量", r"现金及现金等价物",
                                 r"销售商品", r"购买商品", r"收到的税费返还", r"支付给职工", r"取得投资收益",
                                 r"处置固定资产", r"购建固定资产"]
        value_patterns = [r"\d{1,3}(,\d{3})*(\.\d+)?", r"-\d{1,3}(,\d{3})*(\.\d+)?"]
        max_pages = 3  # 现金流量表通常不超过3页

    # 确定结束页面（最多检查max_pages页）
    end_page = min(start_page + max_pages, len(doc) - 1)

    # 直接提取文本并解析
    for page_num in range(start_page, end_page + 1):
        page = doc[page_num]

        # 获取页面文本，按行分割
        text = page.get_text("text")
        lines = text.split('\n')

        # 处理每一行
        for line in lines:
            # 清理行
            line = line.strip()
            if not line:
                continue

            # 检查是否为表格数据行（包含项目名和数值）
            is_data_row = False

            # 检查是否包含项目列关键词
            if any(re.search(pattern, line) for pattern in items_column_patterns):
                is_data_row = True

            # 检查是否包含数值
            if any(re.search(pattern, line) for pattern in value_patterns):
                is_data_row = True

            if is_data_row:
                # 分割行为项目名和数值
                parts = re.split(r'\s{2,}', line)
                if len(parts) >= 2:  # 至少有项目名和一个数值
                    table_data.append(parts)

    return table_data


def clean_and_format_data(extracted_data):
    """清理和格式化提取的数据"""
    cleaned_data = {}

    for table_type, data in extracted_data.items():
        if not data:
            cleaned_data[table_type] = pd.DataFrame()
            continue

        # 处理特定表格
        if table_type == "balance_sheet":
            # 识别和规范资产负债表列名
            df = format_balance_sheet(data)
        elif table_type == "income_statement":
            # 识别和规范利润表列名
            df = format_income_statement(data)
        elif table_type == "cash_flow":
            # 识别和规范现金流量表列名
            df = format_cash_flow_statement(data)

        # 移除空行
        if not df.empty:
            df = df.dropna(how='all')

        cleaned_data[table_type] = df

    return cleaned_data


def format_balance_sheet(data):
    """格式化资产负债表数据"""
    # 创建默认列名
    columns = ["项目", "附注", "期末余额", "期初余额"]

    # 标准化数据行
    formatted_rows = []
    for row in data:
        # 如果行长度少于标准列数，使用空值填充
        if len(row) < len(columns):
            row = row + [''] * (len(columns) - len(row))
        # 如果行长度超过标准列数，截断
        elif len(row) > len(columns):
            row = row[:len(columns)]

        formatted_rows.append(row)

    # 创建DataFrame
    df = pd.DataFrame(formatted_rows, columns=columns)

    # 清理项目列，移除可能的附注号
    df["项目"] = df["项目"].apply(lambda x: re.sub(r'\s*[0-9]+\s*$', '', str(x)))

    # 填充项目列中的空值（向下填充）
    df["项目"] = df["项目"].fillna('')

    return df


def format_income_statement(data):
    """格式化利润表数据"""
    # 创建默认列名
    columns = ["项目", "附注", "本期金额", "上期金额"]

    # 标准化数据行
    formatted_rows = []
    for row in data:
        # 如果行长度少于标准列数，使用空值填充
        if len(row) < len(columns):
            row = row + [''] * (len(columns) - len(row))
        # 如果行长度超过标准列数，截断
        elif len(row) > len(columns):
            row = row[:len(columns)]

        formatted_rows.append(row)

    # 创建DataFrame
    df = pd.DataFrame(formatted_rows, columns=columns)

    # 清理项目列，移除可能的附注号
    df["项目"] = df["项目"].apply(lambda x: re.sub(r'\s*[0-9]+\s*$', '', str(x)))

    # 填充项目列中的空值（向下填充）
    df["项目"] = df["项目"].fillna('')

    return df


def format_cash_flow_statement(data):
    """格式化现金流量表数据"""
    # 创建默认列名
    columns = ["项目", "附注", "本期金额", "上期金额"]

    # 标准化数据行
    formatted_rows = []
    for row in data:
        # 如果行长度少于标准列数，使用空值填充
        if len(row) < len(columns):
            row = row + [''] * (len(columns) - len(row))
        # 如果行长度超过标准列数，截断
        elif len(row) > len(columns):
            row = row[:len(columns)]

        formatted_rows.append(row)

    # 创建DataFrame
    df = pd.DataFrame(formatted_rows, columns=columns)

    # 清理项目列，移除可能的附注号
    df["项目"] = df["项目"].apply(lambda x: re.sub(r'\s*[0-9]+\s*$', '', str(x)))

    # 填充项目列中的空值（向下填充）
    df["项目"] = df["项目"].fillna('')

    return df


def create_formatted_excel(data_dict, output_path):
    """创建一个格式化的Excel文件，简洁易读"""
    wb = Workbook()

    # 删除默认创建的工作表
    default_sheet = wb.active
    wb.remove(default_sheet)

    # 设置样式
    header_fill = PatternFill(start_color="E6E6E6", end_color="E6E6E6", fill_type="solid")
    bold_font = Font(bold=True)
    thin_border = Border(
        left=Side(style='thin'),
        right=Side(style='thin'),
        top=Side(style='thin'),
        bottom=Side(style='thin')
    )
    center_alignment = Alignment(horizontal='center', vertical='center')
    left_alignment = Alignment(horizontal='left', vertical='center')

    # 工作表名称映射
    sheet_names = {
        "balance_sheet": "合并资产负债表",
        "income_statement": "合并利润表",
        "cash_flow": "合并现金流量表"
    }

    # 为每个表创建工作表
    for key, name in sheet_names.items():
        if key in data_dict and not data_dict[key].empty:
            # 创建工作表
            ws = wb.create_sheet(title=name)
            df = data_dict[key]

            # 写入标题行
            for col_idx, col_name in enumerate(df.columns, 1):
                cell = ws.cell(row=1, column=col_idx, value=col_name)
                cell.font = bold_font
                cell.fill = header_fill
                cell.alignment = center_alignment
                cell.border = thin_border

            # 写入数据行
            for row_idx, row in enumerate(df.values, 2):
                # 检查是否为主要项目行（如"流动资产合计"、"资产总计"等）
                is_total_row = False
                item_name = str(row[0])
                if "合计" in item_name or "总计" in item_name:
                    is_total_row = True

                for col_idx, value in enumerate(row, 1):
                    cell = ws.cell(row=row_idx, column=col_idx, value=value)
                    cell.border = thin_border

                    # 设置对齐方式
                    if col_idx == 1:  # 项目列左对齐
                        cell.alignment = left_alignment
                    else:  # 其他列居中
                        cell.alignment = center_alignment

                    # 对合计行应用加粗字体
                    if is_total_row:
                        cell.font = bold_font

            # 调整列宽
            for column in ws.columns:
                col_letter = column[0].column_letter
                max_length = 0
                for cell in column:
                    try:
                        if len(str(cell.value)) > max_length:
                            max_length = len(str(cell.value))
                    except:
                        pass
                adjusted_width = (max_length + 2) * 1.2
                ws.column_dimensions[col_letter].width = adjusted_width

    # 保存Excel文件
    try:
        wb.save(output_path)
        print(f"已将数据保存到: {output_path}")
    except Exception as e:
        print(f"保存Excel文件时出错: {e}")


def main(pdf_path):
    """主函数"""
    # 提取文件名（不含扩展名）作为输出Excel的文件名
    base_name = os.path.splitext(os.path.basename(pdf_path))[0]
    output_dir = os.path.dirname(pdf_path)
    output_path = os.path.join(output_dir, f"{base_name}_财务数据_简洁版.xlsx")

    # 提取财务表格数据
    extracted_data = extract_financial_tables(pdf_path)

    # 检查是否成功提取了数据
    has_data = any(len(data) > 0 for data in extracted_data.values())
    if not has_data:
        print("警告：未能从PDF中提取到任何财务数据。")
        print("可能原因：")
        print("1. PDF文件中不含财务报表")
        print("2. 财务报表格式无法被程序识别")
        print("3. PDF文件结构复杂，无法正确提取文本")
        return None

    # 清理和格式化数据
    formatted_data = clean_and_format_data(extracted_data)

    # 创建格式化的Excel文件
    create_formatted_excel(formatted_data, output_path)

    return output_path


def process_existing_excel(excel_path):
    """处理已经存在的Excel财务数据文件"""
    print(f"正在处理Excel文件: {excel_path}")

    try:
        # 读取Excel文件中的所有表
        xls = pd.ExcelFile(excel_path)
        sheet_names = xls.sheet_names

        # 显示文件中的工作表
        print(f"Excel文件包含以下工作表: {sheet_names}")

        # 读取各个工作表
        data_dict = {}
        for sheet in sheet_names:
            df = pd.read_excel(excel_path, sheet_name=sheet)
            print(f"工作表 '{sheet}' 包含 {len(df)} 行数据")
            data_dict[sheet] = df

            # 显示前5行数据预览
            print(f"'{sheet}' 预览:")
            print(df.head())
            print("\n")

        return data_dict

    except Exception as e:
        print(f"处理Excel文件时出错: {e}")
        return None


if __name__ == "__main__":
    # 两种模式：处理PDF或处理已有Excel
    mode = input("请选择模式 (1: 处理PDF, 2: 处理已有Excel): ")

    if mode == "1":
        # 文件路径
        pdf_path = input("请输入PDF文件路径 (例如: d:\\股票\\宏发股份\\宏发股份 2023年年度报告-20240330.pdf): ")
        if not pdf_path:
            pdf_path = r"d:\股票\宏发股份\宏发股份 2023年年度报告-20240330.pdf"

        # 确保文件存在
        if not os.path.exists(pdf_path):
            print(f"错误：文件 {pdf_path} 不存在")
        else:
            output_file = main(pdf_path)
            if output_file:
                print(f"处理完成，结果保存在: {output_file}")

    elif mode == "2":
        # 处理已有Excel文件
        excel_path = input(
            "请输入Excel文件路径 (例如: d:\\股票\\宏发股份\\宏发股份 2023年年度报告20240330_财务数据_简洁版.xlsx): ")
        if not excel_path:
            excel_path = r"d:\股票\宏发股份\宏发股份 2023年年度报告20240330_财务数据_简洁版.xlsx"

        # 确保文件存在
        if not os.path.exists(excel_path):
            print(f"错误：文件 {excel_path} 不存在")
        else:
            data_dict = process_existing_excel(excel_path)
            if data_dict:
                print("Excel文件处理完成")

                # 提示是否需要进一步分析
                analyze = input("是否需要对数据进行进一步分析? (y/n): ")
                if analyze.lower() == 'y':
                    # 这里可以添加更多的分析功能
                    print("分析功能待开发...")

    else:
        print("无效的选择，程序退出")